package db

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/sha256"
	"database/sql"
	"deepseek-chat/internal/config"
	"deepseek-chat/internal/models"
	"encoding/base64"
	"encoding/hex"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
	"time"

	_ "modernc.org/sqlite"
)

var DB *sql.DB

func Init(cfg *config.Config) error {
	if err := os.MkdirAll(filepath.Dir(cfg.DBPath), 0o755); err != nil {
		return err
	}
	dsn := filepath.ToSlash(fmt.Sprintf("file:%s?_pragma=foreign_keys(ON)", cfg.DBPath))
	var err error
	DB, err = sql.Open("sqlite", dsn)
	if err != nil {
		return err
	}

	return initSchema()
}

func initSchema() error {
	queries := []string{
		`CREATE TABLE IF NOT EXISTS topics (
			id TEXT PRIMARY KEY,
			name TEXT NOT NULL,
			created_at INTEGER NOT NULL
		)`,
		`CREATE TABLE IF NOT EXISTS encrypted_keys (
			id TEXT PRIMARY KEY,
			encrypted_key TEXT NOT NULL,
			salt TEXT NOT NULL,
			created_at INTEGER NOT NULL,
			updated_at INTEGER NOT NULL
		)`,
		`CREATE TABLE IF NOT EXISTS messages (
			id TEXT PRIMARY KEY,
			topic_id TEXT NOT NULL,
			role TEXT NOT NULL,
			content TEXT NOT NULL,
			username TEXT NOT NULL,
			mine BOOLEAN NOT NULL DEFAULT 0,
			timestamp INTEGER NOT NULL,
			created_at INTEGER NOT NULL,
			FOREIGN KEY (topic_id) REFERENCES topics(id) ON DELETE CASCADE
		)`,
		`CREATE INDEX IF NOT EXISTS idx_messages_topic_id ON messages(topic_id)`,
		`CREATE INDEX IF NOT EXISTS idx_messages_timestamp ON messages(timestamp)`,
	}

	for _, query := range queries {
		if _, err := DB.Exec(query); err != nil {
			return err
		}
	}
	return nil
}

func SeedDefaultTopicIfEmpty() error {
	var cnt int
	if err := DB.QueryRow("SELECT COUNT(1) FROM topics").Scan(&cnt); err != nil {
		return err
	}
	if cnt == 0 {
		id, err := RandomID()
		if err != nil { return err }
		now := time.Now().Unix()
		if _, err := DB.Exec("INSERT INTO topics(id, name, created_at) VALUES(?, ?, ?)", id, "默认话题", now); err != nil {
			return err
		}
	}
	return nil
}

func RandomID() (string, error) {
	var b [8]byte
	if _, err := rand.Read(b[:]); err != nil {
		return "", err
	}
	return hex.EncodeToString(b[:]), nil
}

// Encryption Logic

func generateSalt() ([]byte, error) {
	salt := make([]byte, 32)
	_, err := rand.Read(salt)
	return salt, err
}

func deriveKey(password string, salt []byte) []byte {
	hash := sha256.Sum256(append([]byte(password), salt...))
	return hash[:]
}

func encrypt(plaintext, password string) (string, string, error) {
	salt, err := generateSalt()
	if err != nil {
		return "", "", err
	}
	
	key := deriveKey(password, salt)
	
	block, err := aes.NewCipher(key)
	if err != nil {
		return "", "", err
	}
	
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return "", "", err
	}
	
	nonce := make([]byte, gcm.NonceSize())
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		return "", "", err
	}
	
	ciphertext := gcm.Seal(nonce, nonce, []byte(plaintext), nil)
	
	return base64.StdEncoding.EncodeToString(ciphertext), 
	       base64.StdEncoding.EncodeToString(salt), nil
}

func decrypt(ciphertext, password, saltStr string) (string, error) {
	salt, err := base64.StdEncoding.DecodeString(saltStr)
	if err != nil {
		return "", err
	}
	
	key := deriveKey(password, salt)
	
	data, err := base64.StdEncoding.DecodeString(ciphertext)
	if err != nil {
		return "", err
	}
	
	block, err := aes.NewCipher(key)
	if err != nil {
		return "", err
	}
	
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return "", err
	}
	
	nonceSize := gcm.NonceSize()
	if len(data) < nonceSize {
		return "", fmt.Errorf("ciphertext too short")
	}
	
	nonce, ciphertext_bytes := data[:nonceSize], data[nonceSize:]
	plaintext, err := gcm.Open(nil, nonce, ciphertext_bytes, nil)
	if err != nil {
		return "", err
	}
	
	return string(plaintext), nil
}

// getServerPassword returns the master password for encryption.
// It prefers MASTER_KEY env var, falling back to machine ID logic.
func getServerPassword() string {
	if key := os.Getenv("MASTER_KEY"); key != "" {
		return key
	}

	// Fallback logic (same as original for backward compatibility if MASTER_KEY not set)
	hostname, _ := os.Hostname()
	machineID := ""
	if data, err := os.ReadFile("/etc/machine-id"); err == nil {
		machineID = strings.TrimSpace(string(data))
	} else if data, err := os.ReadFile("/var/lib/dbus/machine-id"); err == nil {
		machineID = strings.TrimSpace(string(data))
	}
	
	if machineID == "" {
		machineID = hostname + "-fallback"
	}
	
	return fmt.Sprintf("deepseek-chat-secure-%s", machineID)
}

func SaveEncryptedAPIKey(apiKey string) error {
	password := getServerPassword()
	encryptedKey, salt, err := encrypt(apiKey, password)
	if err != nil {
		return err
	}
	
	now := time.Now().Unix()
	id := "api_key_main"
	
	_, err = DB.Exec(`
		INSERT OR REPLACE INTO encrypted_keys (id, encrypted_key, salt, created_at, updated_at)
		VALUES (?, ?, ?, COALESCE((SELECT created_at FROM encrypted_keys WHERE id = ?), ?), ?)
	`, id, encryptedKey, salt, id, now, now)
	
	return err
}

func LoadDecryptedAPIKey() (string, error) {
	var encKey models.EncryptedKey
	err := DB.QueryRow(`
		SELECT id, encrypted_key, salt, created_at, updated_at 
		FROM encrypted_keys WHERE id = ?
	`, "api_key_main").Scan(&encKey.ID, &encKey.EncryptedKey, &encKey.Salt, &encKey.CreatedAt, &encKey.UpdatedAt)
	
	if err != nil {
		if err == sql.ErrNoRows {
			return "", nil
		}
		return "", err
	}
	
	password := getServerPassword()
	return decrypt(encKey.EncryptedKey, password, encKey.Salt)
}
